import idaapi
import idautils
import idc
import ida_segment
import os
import logging
import coloredlogs
import ida_hexrays

out_folder = idc.ARGV[1]

def set_image_base(new_base):
    # 获取当前图像基地址
    current_base = idaapi.get_imagebase()
    print(f"当前基地址: {hex(current_base)}")
    
    # 计算偏移量
    delta = new_base - current_base
    print(f"新基地址: {hex(new_base)}, 偏移量: {hex(delta)}")
    
    if delta == 0:
        print("基地址相同，无需重新基址。")
        return
    
    # 重新基址程序
    try:
        ida_segment.rebase_program(delta, ida_segment.MSF_FIXONCE)
        print(f"图像基地址已成功设置为 {hex(new_base)}")
    except Exception as e:
        print(f"设置图像基地址时出错: {e}")

def get_function_blocks(func_ea):
    """获取函数的基本块信息"""
    func = idaapi.get_func(func_ea)
    blocks = []
    if func:
        flow_chart = idaapi.FlowChart(func)
        for block in flow_chart:
            blocks.append((block.start_ea, block.end_ea))
    return blocks

def find_block(blocks, address):
    """查找包含特定地址的基本块"""
    for start, end in blocks:
        if start <= address < end:
            return (start, end)
    return (None, None)

def get_function_disasm(func_ea):
    """获取函数的反汇编代码，仅处理.text段内的指�?"""
    # 检查是否在.text段内
    seg_name = idc.get_segm_name(func_ea)
    if seg_name != '.text':
        return []
        
    func_disasm = []
    func = idaapi.get_func(func_ea)
    if not func:
        return []
    
    # 获取函数范围内的指令    
    current_ea = func.start_ea
    while current_ea < func.end_ea:
        # 确保当前地址还在.text段内
        if idc.get_segm_name(current_ea) != '.text':
            break
        disasm_line = idc.generate_disasm_line(current_ea, 0)
        if disasm_line:
            func_disasm.append((current_ea, disasm_line))
        current_ea = idc.next_head(current_ea)
    
    return func_disasm

def get_function_decompilation(func_ea):
    """获取函数的反编译代码，仅处理.text段内的函�?"""
    # 检查是否在.text段内
    seg_name = idc.get_segm_name(func_ea)
    # if seg_name != '.text':
    #     return ""
    
    # func = idaapi.get_func(func_ea)
    # if not func:
    #     return ""
    
    # try:
    #     cfunc = idaapi.decompile(func_ea)
    #     if cfunc:
    #         return str(cfunc)
    # except Exception as e:
    #     print(f"Decompilation error at {hex(func_ea)}: {e}")
    # return ""
    if seg_name == '.text':
        func_name = idc.get_func_name(func_ea)
        try:
            cfunc = idaapi.decompile(func_ea)
            if cfunc:
                return str(cfunc)
        except idaapi.DecompilationFailure:
            return f"// Failed to decompile {func_name} at {hex(func_ea)}\n"
    return ""

def detect_inlining_per_function(func_ea, func_name, func_disasm, func_decomp, inline_disasm_segments):
    """检测特定敏感函数的内联情况"""
    # 定义需要检测的敏感函数列表
    sensitive_functions = [
        'strcpy', 'memcpy', 'memmove',  # 基本的内存操作函�?
        'strcat', 'sprintf', 'gets',    # 其他可能导致缓冲区溢出的函数
        'wcscpy', 'wcscat' , 'memset' , # 宽字符相关的危险函数
        'fgets', 'recv',
        'strncat', 'strncpy',
    ]
    
    # 构建匹配模式，只匹配特定的敏感函�?
    pattern = r'\b(' + '|'.join(sensitive_functions) + r')\s*\([^)]*\)'
    function_calls = list(re.finditer(pattern, func_decomp))

    for call in function_calls:
        function_name = call.group(1)
        decomp_call_offset = call.start()
        found_call = False
        basic_block_start = None
        inline_decompile = ""
        instructions = ""
        last_instructions = ""
        flag = 0
        decompile = []
        call_instructions = []

        # 在汇编代码中查找对应的call指令
        for address, asm_line in func_disasm:
            # 匹配函数调用（允许有下划线前缀�?
            pattern = rf'\bcall\s+(_)?{re.escape(function_name)}\b'
            if re.search(pattern, asm_line):
                found_call = True
                # 获取包含call指令的基本块
                func = idaapi.get_func(func_disasm[0][0])
                if func:
                    for block in idaapi.FlowChart(func):
                        if block.start_ea <= address < block.end_ea:
                            basic_block_start = block.start_ea
                            break                
                break
   
        # 如果没有找到call指令，认为是内联
        if not found_call:
            # 使用函数的第一个地址作为参考点
            basic_block_start = func_disasm[0][0] if func_disasm else None
            # 输出此时的反编译代码
            for line_num, line in enumerate(func_decomp.splitlines()):
                if function_name in line:
                    if line_num > 0:
                        previous_line = func_decomp.splitlines()[line_num - 1]
                        previous_line = previous_line.lstrip()
                        # inline_decompile += previous_line.lstrip() + "\n"
                    # 去除开头的空格
                    inline_decomp = "\n".join(correct_line.lstrip() for correct_line in line.splitlines())
                    inline_decompile += inline_decomp
                    # print(repr(inline_decomp))
                    # print(inline_decomp)
                    last_asm_lines = get_asm_from_decompiled_line(func_ea, previous_line)
                    asm_lines = get_asm_from_decompiled_line(func_ea, inline_decomp)
                    if asm_lines and last_asm_lines:
                        instructions += asm_lines[-1]
                        last_instructions += last_asm_lines[-1]     
 
                        # flag += 1
        # 记录内联信息（无论是否内联都记录，方便后续分析）
        inline_disasm_segments.append((
            func_name,           # 当前分析的函数名
            function_name,       # 检测到的敏感函数名
            func_disasm,         # 反汇编代码
            decomp_call_offset,  # 在反编译代码中的偏移
            basic_block_start,   # 基本块起始地址
            not found_call,      # 是否被内联
            inline_decompile,     # 内联的反编译代码
            instructions,         # 内联的行汇编代码
            last_instructions,   # 上一行的行汇编代码
            flag    # 计数标志
        ))                      
        # else:
        #     basic_block_start = func_disasm[0][0] if func_disasm else None
        #     # 输出此时的反编译代码
        #     for line_num, line in enumerate(func_decomp.splitlines()):
        #         if function_name in line and any(f"{func}(" in line for func in sensitive_functions):
        #             # 去除开头的空格
        #             inline_decomp = "\n".join(correct_line.lstrip() for correct_line in line.splitlines())
        #             decompile.append(inline_decomp)
        #             # print(repr(inline_decomp))
        #             print(inline_decomp)
        #             print(decompile) 
        #             asm_lines = get_asm_from_decompiled_line(func_ea, inline_decomp)
        #             print(asm_lines)
        #             if asm_lines:
        #                 call_instructions.append(asm_lines[-1])
        #                 print(call_instructions)    
                    
        #     inline_disasm_segments.append((
        #         func_name,           # 当前分析的函数名
        #         function_name,       # 检测到的敏感函数名
        #         func_disasm,         # 反汇编代码
        #         decomp_call_offset,  # 在反编译代码中的偏移
        #         basic_block_start,   # 基本块起始地址
        #         not found_call,      # 是否被内联
        #         decompile,     # 内联的反编译代码
        #         call_instructions,         # 内联的行汇编代码
        #         last_instructions,   # 上一行的行汇编代码
        #         flag    # 计数标志
        #     ))    
                            

        # print(flag)
       

# def get_decompiled_line_from_instruction(line, ea):
#     if not ida_hexrays.init_hexrays_plugin():
#         print("Hex-Rays decompiler is not available.")
#         return ""
    
#     func = idaapi.get_func(ea)
#     if not func:
#         print(f"No function found at address 0x{ea:X}")
#         return ""
    
#     cfunc = idaapi.decompile(func)
#     if not cfunc:
#         print(f"Failed to decompile function at address 0x{ea:X}")
#         return ""
    
#     # 提取函数名，例如从 'call    _strcpy' 中提取 '_strcpy'
#     match = re.search(r'call\s+([\w@]+)', line)
#     if match:
#         called_func = match.group(1).lstrip('_')  # 去掉前导下划线
#     else:
#         print("Failed to extract function name from line")
#         return ""
    
#     # 获取反编译后的伪代码行
#     sv = cfunc.get_pseudocode()
#     if not sv:
#         print("Failed to get pseudocode")
#         return ""
    
#     # 将伪代码行转换为字符串列表，并移除格式化标签
#     code_lines = [ida_lines.tag_remove(sv[i].line) for i in range(sv.size())]
    
#     # 在反编译代码中搜索函数调用
#     matched_line = ""
#     for code_line in code_lines:
#         # 检查代码行中是否包含函数调用
#         if re.search(r'\b' + re.escape(called_func) + r'\b', code_line):
#             matched_line = code_line.strip()
#             break
    
#     return matched_line

def get_decompiled_line_from_instruction(line, ea):
    class CallFinder(idaapi.ctree_visitor_t):
        def __init__(self, ea):
            idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST)
            self.ea = ea
            self.call_item = None
            self.found = False

        def visit_expr(self, expr):
            if expr.op == idaapi.cot_call and expr.ea == self.ea:
                self.call_item = expr
                self.found = True
                return 1
            return 0
    if not ida_hexrays.init_hexrays_plugin():
        print("Hex-Rays decompiler is not available.")
        return ""
    
    func = idaapi.get_func(ea)
    if not func:
        print(f"No function found at address 0x{ea:X}")
        return ""
    
    cfunc = idaapi.decompile(func)
    if not cfunc:
        print(f"Failed to decompile function at address 0x{ea:X}")
        return ""

    # 使用CallFinder查找特定调用
    finder = CallFinder(ea)
    finder.apply_to(cfunc.body, None)

    if finder.found and finder.call_item:
        try:
            # 获取调用表达式
            call_item = finder.call_item
            
            # 获取参数信息
            args = []
            if call_item.a:
                for i in range(call_item.a.size()):
                    arg = call_item.a[i]
                    args.append(str(arg))

            # 获取对应的反编译代码行
            item_line = ida_lines.tag_remove(call_item.print1(None)).strip()
            if item_line:
                return item_line

            # 如果上面方法失败，尝试遍历所有代码行
            sv = cfunc.get_pseudocode()
            for i in range(len(sv)):
                line_item = sv[i]
                if ea in [item.ea for item in cfunc.treeitems]:
                    return ida_lines.tag_remove(line_item.line).strip()
        except Exception as e:
            print(f"Error processing call at 0x{ea:X}: {str(e)}")

    # 回退到原有的模式
    match = re.search(r'call\s+([\w@]+)', line)
    if match:
        called_func = match.group(1).lstrip('_')
        sv = cfunc.get_pseudocode()
        if sv:
            code_lines = [ida_lines.tag_remove(sv[i].line) for i in range(sv.size())]
            for code_line in code_lines:
                if re.search(r'\b' + re.escape(called_func) + r'\b', code_line):
                    return code_line.strip()
    
    return ""


def get_asm_from_decompiled_line(ea, line_text):
    if not ida_hexrays.init_hexrays_plugin():
        print("Hex-Rays decompiler is not available.")
        return []

    cfunc = idaapi.decompile(ea)
    if not cfunc:
        print("Failed to decompile function at 0x%X" % ea)
        return []

    class ctree_visitor(idaapi.ctree_visitor_t):
        def __init__(self, *args):
            idaapi.ctree_visitor_t.__init__(self, *args)
            self.found = False
            self.ea_set = set()

        def visit_insn(self, i):
            line = idaapi.tag_remove(i.print1(None))
            if line_text in line:
                if i.ea != idaapi.BADADDR:
                    self.ea_set.add(i.ea)
                self.found = True
            return 0

    visitor = ctree_visitor(idaapi.CV_FAST)
    visitor.apply_to(cfunc.body, None)

    if visitor.found:
        asm_instructions = []
        for addr in sorted(visitor.ea_set):
            disasm = idc.GetDisasm(addr)
            asm_instructions.append("0x%X: %s" % (addr, disasm))
        return asm_instructions
    else:
        print("未找到匹配的反编译代码行。")
        return []

def main():
    sensitive_functions = [
        'strcpy', 'memcpy', 'memmove',  # 基本的内存操作函�?
        'strcat', 'sprintf', 'gets',    # 其他可能导致缓冲区溢出的函数
        'wcscpy', 'wcscat' , 'memset', # 宽字符相关的危险函数
        'fgets', 'recv',
        'strncat', 'strncpy',
    ]    
    
  # 获取当前图像基地址
    current_base = idaapi.get_imagebase()
    print(f"当前基地址: {hex(current_base)}")

    # 定义新的基地址
    desired_base = 0x400000  # 设置您希望的基地址

    # 检测是否需要设置新的基地址
    if current_base == 0x0:
        print("检测到当前基地址为 0x0，准备设置新的基地址。")
        set_image_base(desired_base)
    else:
        print(f"当前基地址已设置为 {hex(current_base)}，无需重新基址。")
        
    # 获取二进制文件名（不含扩展名�?
    binary_name = os.path.splitext(os.path.basename(idaapi.get_input_file_path()))[0]
    binary_path = idaapi.get_input_file_path()
    
    # 创建结果目录
    # result_dir = os.path.join(os.path.dirname(binary_path), 'result')
    # os.makedirs(result_dir, exist_ok=True)
    
    result_dir = out_folder
    
    # 输出文件路径
    output_disasm_path = f"{result_dir}/{binary_name}_assembly_code.txt"
    output_decomp_path = f"{result_dir}/{binary_name}_decompiled_code.c"
    output_inline_disasm_path = f"{result_dir}/{binary_name}_inline_disassembly.txt"
    # block_info_path = f"{result_dir}/{binary_name}_block_info.txt"
    
    with open(output_disasm_path, 'w') as disasm_file, \
         open(output_decomp_path, 'w') as decomp_file:
        #  open(block_info_path, 'w') as block_file
        
        inline_disasm_segments = []
        
        # 遍历所有函数
        for func_ea in idautils.Functions():
            # 检查函数是否在.text段内
            if idc.get_segm_name(func_ea) != '.text':
                continue
                
            func_name = idc.get_func_name(func_ea)
            # if not func_name or func_name.startswith("sub_"):
            #     continue

            func_disasm = get_function_disasm(func_ea)
            if not func_disasm:
                continue

            func_decomp = get_function_decompilation(func_ea)
            if not func_decomp:
                continue

            # 保存反汇编代码
            disasm_file.write(f"\n\n; Function: {func_name}\n")
            disasm_file.write("\n".join(f"{hex(address)}: {line}" for address, line in func_disasm) + "\n")
            
            # 保存反编译代码
            decomp_file.write(f"\n\n// Function: {func_name}\n")
            decomp_file.write(func_decomp + "\n")
            
            # 获取并保存基本块信息
            # blocks = get_function_blocks(func_ea)
            # block_file.write(f"\n// Function: {func_name}\n")
            # for start, end in blocks:
            #     block_file.write(f"Block Start: {hex(start)}, Block End: {hex(end)}\n")
            
            # 检测内�?
            detect_inlining_per_function(func_ea, func_name, func_disasm, func_decomp, inline_disasm_segments)


    # 保存内联检测结�?
    with open(output_inline_disasm_path, 'w') as inline_file:
        # 在文件顶部定义一个集合来存储已处理的调用地址
        processed_calls = set()
        processed_func = set()
        # print(inline_disasm_segments)  
        for func_name, inlined_func_name, asm_lines, decomp_offset, basic_block_start, is_inlined, inline_decomp, instructions, last_instructions, flag in inline_disasm_segments:
            # print(inlined_func_name)
            if is_inlined:
                # # 获取地址
                # last_address = int(last_instructions.split(':')[0], 16)
                # current_address = int(instructions.split(':')[0], 16)
            
                # # 找到相应位置
                # start_index = next((i for i, (address, _) in enumerate(asm_lines) if address == last_address), None)
                # 获取地址
                current_address = int(instructions.split(':')[0], 16)
                if last_instructions:
                    last_address = int(last_instructions.split(':')[0], 16)
                    # 找到相应位置
                    start_index = next((i for i, (address, _) in enumerate(asm_lines) if address == last_address), None)
                else:
                    # 如果没有上一行的汇编代码，使用函数的起始地址
                    start_index = 0            
                
                end_index = next((i for i, (address, _) in enumerate(asm_lines) if address == current_address), None)
                # inline_decomp_str = inline_decomp_str.replace("\\", "\\\\").replace("\"", "\\\"")
                if start_index is not None and end_index is not None:
                    # 输出相应位置及其中间的代码
                    inline_file.write(f"\n\n// Function '{inlined_func_name}' in '{func_name}' - INLINED:\n")
                    inline_file.write(f"// Occurrence in decompiled code at offset {decomp_offset}\n")
                    inline_file.write(f"Inline Decompile Code '{inline_decomp}'\n")
                    inline_file.write(f"'{instructions}'\n")
                    inline_file.write(f"// Near address: {hex(basic_block_start) if basic_block_start else 'N/A'}\n")
                    inline_file.write("\n".join(f"{hex(address)}: {line}" for address, line in asm_lines[start_index:end_index+1]) + "\n")
                else:
                    inline_file.write(f"\n\n// Function '{inlined_func_name}' in '{func_name}' - INLINED:\n")
                    inline_file.write(f"// Occurrence in decompiled code at offset {decomp_offset}\n")
                    inline_file.write(f"Inline Decompile Code '{inline_decomp}'\n")
                    inline_file.write(f"'{instructions}'\n")
                    inline_file.write(f"// Near address: {hex(basic_block_start) if basic_block_start else 'N/A'}\n")
                    inline_file.write("\n".join(f"{hex(address)}: {line}" for address, line in asm_lines) + "\n")
                # inline_file.write(f"\n\n// Function '{inlined_func_name}' in '{func_name}' - INLINED:\n")
                # inline_file.write(f"// Occurrence in decompiled code at offset {decomp_offset}\n")
                # inline_file.write(f"Inline Decomplie Code '{inline_decomp}'\n")
                # inline_file.write(f"'{instructions}'\n")
                # last_instructions
                # instructions
                # inline_file.write(f"// Near address: {hex(basic_block_start) if basic_block_start else 'N/A'}\n")
                # inline_file.write("\n".join(f"{hex(address)}: {line}" for address, line in asm_lines) + "\n")
            else:
                # print(inlined_func_name)
                # 获取地址
                # current_address = int(instructions.split(':')[0], 16)
                # call_addresses = [hex(address) for address, line in asm_lines if 'call' in line]

                # inline_file.write(f"\n// Function '{inlined_func_name}' in '{func_name}' - NOT INLINED\n")
                #     inline_file.write(f"'{instructions[i]}'\n")
                # inline_file.write(f"// Call addresses: {', '.join(call_addresses)}\n")
                #     inline_file.write(f"Decompile Code '{inline_decomp[i]}'\n")
                print(asm_lines)
                for address, line in asm_lines:
                    # print(inlined_func_name)
                    if not is_inlined and any(f"{func}" in line for func in sensitive_functions):
                        print(inlined_func_name,address,"1")
                        # 获取调用地址（假设instructions包含地址信息）                     
                        # 检查是否已经处理过这个调用
                        if address not in processed_calls and inlined_func_name not in processed_func:
                            processed_calls.add(address)
                            processed_func.add(inlined_func_name)
                            print(inlined_func_name,address,"2")
                            inline_file.write(f"\n// Function '{inlined_func_name}' in '{func_name}' - NOT INLINED\n")
                            inline_file.write(f"{hex(address)}: {line}\n")
                            decomp_line = get_decompiled_line_from_instruction(line, address)
                            if decomp_line:
                                inline_file.write(f"Decompiled Line: {decomp_line}\n")
                            else:
                                inline_file.write(f"Decompiled Line: N/A\n")     

    print(f"Analysis complete. Results saved in {result_dir}")

if __name__ == "__main__":
    idaapi.autoWait()
    main()
    idc.qexit(0)
